// AH-FastPkt - 一个用于发送HTTP请求并检测WAF拦截的工具
package main

import (
	"bufio"
	"crypto/tls"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"

	"github.com/sirupsen/logrus"
)

// init 函数在程序启动时执行，用于初始化配置
func init() {
	// 设置命令行参数的使用说明
	flag.Usage = func() {
		fmt.Printf("用法: %s --target <目标url> --filename <脚本路径> [--status <403>] [--echo_res <false/true>] [--echo_res_code <false/true>]\n", os.Args[0])
		flag.PrintDefaults()
	}
	// 获取当前工作目录
	currentDir, err := os.Getwd()
	if err != nil {
		fmt.Printf("err: %v\n", err)
	}
	// 创建日志目录
	logDir := filepath.Join(currentDir, "log")

	// 检查日志目录是否存在，不存在则创建
	if _, err := os.Stat(logDir); os.IsNotExist(err) {
		err := os.Mkdir(logDir, os.ModePerm)
		if err != nil {
			fmt.Printf("err: %v\n", err)
		}
	}
}

func main() {
	// 定义命令行参数
	target := flag.String("target", "", "指定目标url,格式http/https://domain:port,记得删除最后一个/")
	filename := flag.String("filename", "", "指定脚本文件的路径")
	status := flag.Int("status", 403, "指定状态码,默认403")
	echo_res_code := flag.Bool("echo_res_code", false, "是否打印状态码,默认不打印")
	echo_res := flag.Bool("echo_res", false, "是否打印请求内容和响应内容,默认不打印")
	help := flag.Bool("help", false, "显示帮助信息")

	flag.Parse()

	// 检查必要参数
	if *help {
		flag.Usage()
		return
	}
	if *target == "" {
		fmt.Println("请提供 --target 参数")
		flag.Usage()
		return
	}

	if *filename == "" {
		fmt.Println("请提供 --filename 参数")
		flag.Usage()
		return
	}

	// 初始化计数器变量
	var code_waf int   // WAF拦截计数
	var code_2xx int   // 2xx状态码计数
	var block int      // 是否被拦截标志
	var code_4xx int   // 4xx状态码计数
	var code_5xx int   // 5xx状态码计数
	var code_other int // 其他状态码计数

	// 初始化日志记录器
	var log = logrus.New()
	var Errorlog = logrus.New()

	// 配置日志格式为JSON
	log.Formatter = new(logrus.JSONFormatter)
	Errorlog.Formatter = new(logrus.JSONFormatter)

	// 配置日志字段映射
	log.SetFormatter(&logrus.JSONFormatter{
		FieldMap: logrus.FieldMap{
			logrus.FieldKeyTime:  "@timestamp",
			logrus.FieldKeyLevel: "@level",
			logrus.FieldKeyMsg:   "@message",
			logrus.FieldKeyFunc:  "@caller",
			"Target":             "@Target",
			"filename":           "@filename",
			"Block":              "@Block",
			"StatusCode":         "@StatusCode",
		},
	})

	Errorlog.SetFormatter(&logrus.JSONFormatter{
		FieldMap: logrus.FieldMap{
			logrus.FieldKeyTime:  "@timestamp",
			logrus.FieldKeyLevel: "@level",
			logrus.FieldKeyMsg:   "@message",
			logrus.FieldKeyFunc:  "@caller",
			"Target":             "@Target",
			"filename":           "@filename",
			"Error_Type":         "@Error_Type",
		},
	})

	// 打开日志文件
	log_file, logerr := os.OpenFile("./log/info.log", os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
	errorlog_file, errlogerr := os.OpenFile("./log/error.log", os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
	if logerr == nil && errlogerr == nil {
		log.Out = log_file
		Errorlog.Out = errorlog_file
	}
	defer log_file.Close()
	defer errorlog_file.Close()

	// 发送HTTP请求并获取状态码
	StatusCode := httpsend(*filename, *target, Errorlog, *echo_res)
	if *echo_res_code {
		fmt.Println(StatusCode)
	}

	// 判断是否被WAF拦截
	if StatusCode != *status {
		block = 0
	} else {
		block = 1
	}

	// 记录WAF拦截日志
	log.WithFields(logrus.Fields{
		"Target":     *target,
		"Block":      block,
		"filename":   filename,
		"StatusCode": StatusCode,
	}).Info("WAF拦截记录")

	// 统计不同状态码的数量
	if StatusCode == *status {
		code_waf += 1
	} else if 200 <= StatusCode && StatusCode <= 299 {
		code_2xx += 1
	} else if 400 <= StatusCode && StatusCode <= 499 {
		code_4xx += 1
	} else if 500 <= StatusCode && StatusCode <= 599 {
		code_5xx += 1
	} else {
		code_other += 1
	}
}

// httpsend 函数用于发送HTTP请求并处理响应
// 参数:
//   - file_path: HTTP请求文件的路径
//   - target: 目标URL
//   - Errorlog: 错误日志记录器
//   - echo_res: 是否打印响应内容
//
// 返回:
//   - statuscode: HTTP响应状态码
func httpsend(file_path string, target string, Errorlog *logrus.Logger, echo_res bool) (statuscode int) {
	// 打开请求文件
	file, err := os.Open(file_path)
	if err != nil {
		Errorlog.WithFields(logrus.Fields{
			"Target":     target,
			"filename":   file_path,
			"Error_Type": "文件错误",
		}).Error(err)
		return
	}
	defer file.Close()

	// 初始化HTTP请求相关变量
	var method string
	var url string
	requestHeaders := make(map[string]string)
	var requestBody string
	readingBody := false

	// 逐行解析请求文件内容
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := scanner.Text()
		if readingBody {
			// 读取请求正文
			requestBody += "\n" + line + "\n"
		} else if line == "" {
			// 空行表示请求头结束，开始读取请求正文
			readingBody = true
		} else if method == "" && url == "" {
			// 解析请求方法和URL
			parts := strings.Fields(line)
			if len(parts) == 3 {
				method = parts[0]
				url = parts[1]
			} else if len(parts) > 3 {
				method = parts[0]
				parts_url := parts[1 : len(parts)-1]
				result := ""
				for _, element := range parts_url {
					result += element + " "
				}
				url = result
			}
		} else {
			// 解析请求头
			parts := strings.SplitN(line, ":", 2)
			if len(parts) == 2 {
				key := strings.TrimSpace(parts[0])
				value := strings.TrimSpace(parts[1])
				requestHeaders[key] = value
			}
		}
	}

	// 检查文件读取错误
	if err := scanner.Err(); err != nil {
		Errorlog.WithFields(logrus.Fields{
			"Target":     target,
			"filename":   file_path,
			"Error_Type": "读文件错误",
		}).Error(err)
		return
	}

	// 创建HTTP请求
	fullURL := target + url
	request, err := http.NewRequest(method, fullURL, strings.NewReader(requestBody))

	if err != nil {
		Errorlog.WithFields(logrus.Fields{
			"Target":     target,
			"filename":   file_path,
			"Error_Type": "创建请求失败",
		}).Error(err)
		return
	}
	request.Header = http.Header{}

	// 设置请求头
	for key, value := range requestHeaders {
		request.Header.Set(key, value)
	}

	// 配置HTTP客户端并发送请求
	client := &http.Client{
		Transport: &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, // 允许不安全的HTTPS连接
		},
	}
	response, err := client.Do(request)
	if err != nil {
		Errorlog.WithFields(logrus.Fields{
			"Target":     target,
			"filename":   file_path,
			"Error_Type": "发送请求失败",
		}).Error(err)
		return
	}
	defer response.Body.Close()

	// 如果需要，打印请求和响应详情
	if echo_res {
		fmt.Printf("file_name: %v\n", file_path)
		fmt.Println()
		fmt.Println("HTTP request:")
		fmt.Printf("%s %s HTTP/1.1\n", method, url)
		for key, values := range request.Header {
			for _, value := range values {
				fmt.Printf("%s: %s\n", key, value)
			}
		}
		fmt.Println(requestBody)
		fmt.Println()
		fmt.Println("HTTP response:")
		// 打印响应状态行
		fmt.Printf("HTTP/%d.%d %s\n", response.ProtoMajor, response.ProtoMinor, response.Status)
		// 打印响应头
		for key, values := range response.Header {
			for _, value := range values {
				fmt.Printf("%s: %s\n", key, value)
			}
		}
		// 打印响应体
		body, err := io.ReadAll(response.Body)
		fmt.Println("响应信息:")
		if err != nil {
			Errorlog.WithFields(logrus.Fields{
				"Target":     target,
				"filename":   file_path,
				"Error_Type": "读取响应失败",
			}).Error(err)
		}
		fmt.Println(string(body))
	}
	return response.StatusCode
}
